"""
多Agent汇率预测分析系统
结合汇率预测模型和舆论政策事件分析，使用LangGraph框架实现多Agent协作
"""

import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datetime import datetime, timedelta
from typing import TypedDict, Dict, Any, Optional, List
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error, r2_score
from langgraph.graph import StateGraph, END
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

# 创建输出目录
output_dir = 'output_images'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# --- 1. 定义多Agent系统的状态 ---
class MultiAgentState(TypedDict):
    # 基础数据
    currency_pair: str
    start_date: str
    end_date: str
    raw_data: Optional[pd.DataFrame]
    
    # 汇率预测相关
    model: Optional[LinearRegression]
    scaler: Optional[StandardScaler]
    predictions: Optional[np.ndarray]
    prediction_dates: Optional[pd.DatetimeIndex]
    plot_path: Optional[str]
    
    # 舆论政策事件分析相关
    sentiment_data: Optional[pd.DataFrame]
    policy_data: Optional[pd.DataFrame]
    event_data: Optional[pd.DataFrame]
    policy_events: Optional[Dict]
    major_events: Optional[Dict]
    combined_analysis: Optional[pd.DataFrame]
    analysis_results: Optional[Dict]
    
    # 用户交互
    user_query: Optional[str]
    prediction_days: Optional[int]
    
    # 系统状态
    next_action: str
    error_message: Optional[str]
    current_agent: str
    agent_results: Dict[str, Any]

# --- 2. 数据处理和模型训练函数 ---
def load_exchange_rate_data(file_path: str = None, currency_pair: str = 'USD/CNY', 
                            start_date: str = '2023-01-01', end_date: str = '2023-12-31') -> pd.DataFrame:
    """加载汇率数据"""
    try:
        if file_path and os.path.exists(file_path):
            print(f"读取真实汇率数据: {file_path}")
            df = pd.read_csv(file_path, skiprows=3)
            df.columns = ['Date', 'exchange_rate'] if len(df.columns) == 2 else ['Date', 'exchange_rate'] + [f'Col_{i}' for i in range(len(df.columns)-2)]
            df = df.dropna(subset=['Date', 'exchange_rate'])
            df['Date'] = pd.to_datetime(df['Date'], errors='coerce')
            df = df.dropna(subset=['Date'])
            
            start = pd.to_datetime(start_date)
            end = pd.to_datetime(end_date)
            df = df[(df['Date'] >= start) & (df['Date'] <= end)]
            df.set_index('Date', inplace=True)
            df['exchange_rate'] = pd.to_numeric(df['exchange_rate'], errors='coerce')
            df = df.sort_index(ascending=True)
            
            if df['exchange_rate'].isnull().any():
                df['exchange_rate'] = df['exchange_rate'].fillna(method='ffill')
                df['exchange_rate'] = df['exchange_rate'].fillna(method='bfill')
            
            if not df.empty:
                print(f"成功加载{currency_pair}真实汇率数据，共{len(df)}条记录")
                return df
    except Exception as e:
        print(f"加载真实汇率数据失败: {str(e)}，使用模拟数据代替")
    
    # 生成模拟数据
    print("生成模拟汇率数据...")
    date_range = pd.date_range(start=start_date, end=end_date)
    
    if currency_pair == 'USD/CNY':
        base_rate = np.linspace(6.8, 7.2, len(date_range))
    elif currency_pair == 'EUR/USD':
        base_rate = np.linspace(1.05, 1.15, len(date_range))
    else:  # EUR/CNY
        base_rate = np.linspace(7.5, 8.0, len(date_range))
    
    seasonal = 0.1 * np.sin(np.linspace(0, 4*np.pi, len(date_range)))
    random_fluctuation = np.random.normal(0, 0.03, len(date_range))
    exchange_rate = base_rate + seasonal + random_fluctuation
    
    df = pd.DataFrame({
        'exchange_rate': exchange_rate
    }, index=date_range)
    
    return df

def create_features(data: pd.DataFrame) -> tuple:
    """创建特征"""
    data = data.copy()
    data['ma_5'] = data['exchange_rate'].rolling(window=5).mean()
    data['ma_10'] = data['exchange_rate'].rolling(window=10).mean()
    data['volatility'] = data['exchange_rate'].rolling(window=5).std()
    data['price_change'] = data['exchange_rate'].pct_change()
    data['lag_1'] = data['exchange_rate'].shift(1)
    data['lag_2'] = data['exchange_rate'].shift(2)
    
    data = data.dropna()
    
    feature_columns = ['ma_5', 'ma_10', 'volatility', 'price_change', 'lag_1', 'lag_2']
    X = data[feature_columns]
    y = data['exchange_rate']
    
    return X, y, data.index

def train_model(X: pd.DataFrame, y: pd.Series) -> tuple:
    """训练模型"""
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    model = LinearRegression()
    model.fit(X_scaled, y)
    
    y_pred = model.predict(X_scaled)
    mse = mean_squared_error(y, y_pred)
    r2 = r2_score(y, y_pred)
    
    print(f"模型训练完成 - MSE: {mse:.6f}, R²: {r2:.4f}")
    
    return model, scaler

def predict_future(model: LinearRegression, scaler: StandardScaler, 
                  last_data: pd.DataFrame, days: int = 30) -> tuple:
    """预测未来汇率"""
    predictions = []
    current_data = last_data.copy()
    
    last_date = current_data.index[-1]
    prediction_dates = pd.date_range(start=last_date + timedelta(days=1), periods=days)
    
    for i in range(days):
        # 创建特征
        current_data['ma_5'] = current_data['exchange_rate'].rolling(window=5).mean()
        current_data['ma_10'] = current_data['exchange_rate'].rolling(window=10).mean()
        current_data['volatility'] = current_data['exchange_rate'].rolling(window=5).std()
        current_data['price_change'] = current_data['exchange_rate'].pct_change()
        current_data['lag_1'] = current_data['exchange_rate'].shift(1)
        current_data['lag_2'] = current_data['exchange_rate'].shift(2)
        
        # 获取最新特征
        latest_features = current_data[['ma_5', 'ma_10', 'volatility', 'price_change', 'lag_1', 'lag_2']].iloc[-1:]
        latest_features = latest_features.dropna()
        
        if latest_features.empty:
            break
            
        # 预测
        features_scaled = scaler.transform(latest_features)
        pred = model.predict(features_scaled)[0]
        predictions.append(pred)
        
        # 更新数据
        new_date = prediction_dates[i]
        current_data.loc[new_date] = {'exchange_rate': pred}
    
    return np.array(predictions), prediction_dates

def plot_forecast(data: pd.DataFrame, prediction_dates: pd.DatetimeIndex, 
                 predictions: np.ndarray, currency_pair: str = 'USD/CNY') -> str:
    """绘制预测图表"""
    plt.figure(figsize=(15, 8))
    
    # 绘制历史数据
    plt.plot(data.index, data['exchange_rate'], 'b-', label='历史汇率', linewidth=2)
    
    # 绘制预测数据 - 使用更明显的红色虚线
    plt.plot(prediction_dates, predictions, 'r--', label='预测汇率', linewidth=4, alpha=0.9, markersize=8)
    
    # 连接历史和预测数据
    if len(data) > 0 and len(predictions) > 0:
        plt.plot([data.index[-1], prediction_dates[0]], 
                [data['exchange_rate'].iloc[-1], predictions[0]], 'g-', alpha=0.5)
    
    plt.title(f'{currency_pair} 汇率预测', fontsize=16)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('汇率', fontsize=12)
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # 根据实际数据范围设置x轴
    if len(data) > 0 and len(prediction_dates) > 0:
        # 获取数据的实际时间范围
        data_start = data.index[0]
        data_end = prediction_dates[-1]
        
        # 设置x轴范围，留一些边距
        time_span = data_end - data_start
        margin = time_span * 0.05  # 5%的边距
        plt.xlim(data_start - margin, data_end + margin)
        
        # 根据时间跨度设置合适的刻度
        if time_span.days > 365:  # 超过一年，按月显示
            plt.gca().xaxis.set_major_locator(mdates.MonthLocator(interval=2))
            plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y年%m月'))
            plt.gca().xaxis.set_minor_locator(mdates.MonthLocator())
        elif time_span.days > 90:  # 超过三个月，按月显示
            plt.gca().xaxis.set_major_locator(mdates.MonthLocator())
            plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y年%m月'))
            plt.gca().xaxis.set_minor_locator(mdates.WeekdayLocator())
        else:  # 短期数据，按周显示
            plt.gca().xaxis.set_major_locator(mdates.WeekdayLocator())
            plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%m月%d日'))
            plt.gca().xaxis.set_minor_locator(mdates.DayLocator())
    
    plt.gcf().autofmt_xdate()
    
    # 保存图表
    filename = f"{currency_pair.replace('/', '_')}_forecast.png"
    filepath = os.path.join(output_dir, filename)
    plt.savefig(filepath, dpi=300, bbox_inches='tight')
    plt.close()
    
    return filepath

# --- 3. 舆论政策事件分析函数 ---
def generate_sentiment_data(start_date: str, end_date: str, exchange_rate_data: pd.DataFrame = None) -> pd.DataFrame:
    """生成与汇率相关的舆论情绪数据"""
    date_range = pd.date_range(start=start_date, end=end_date)
    
    if exchange_rate_data is not None and len(exchange_rate_data) > 0:
        # 基于汇率变化生成相关的情绪数据
        rate_changes = exchange_rate_data['exchange_rate'].pct_change().fillna(0)
        
        # 情绪与汇率变化负相关（汇率上升时情绪下降）
        base_sentiment = 0.5 - 0.3 * rate_changes.rolling(window=5).mean().fillna(0)
        
        # 添加一些滞后效应
        lagged_effect = -0.2 * rate_changes.shift(1).fillna(0)
        
        # 添加随机波动
        random_sentiment = np.random.normal(0, 0.05, len(base_sentiment))
        
        sentiment = base_sentiment + lagged_effect + random_sentiment
        
        # 确保在合理范围内
        sentiment = np.clip(sentiment, 0.1, 0.9)
        
        sentiment_df = pd.DataFrame({
            'sentiment': sentiment
        }, index=exchange_rate_data.index)
    else:
        # 如果没有汇率数据，生成基础情绪数据
        base_sentiment = 0.5 + 0.2 * np.sin(np.linspace(0, 4*np.pi, len(date_range)))
        random_sentiment = np.random.normal(0, 0.1, len(date_range))
        sentiment = np.clip(base_sentiment + random_sentiment, 0, 1)
        
        sentiment_df = pd.DataFrame({
            'sentiment': sentiment
        }, index=date_range)
    
    return sentiment_df

def generate_policy_data(start_date: str, end_date: str, exchange_rate_data: pd.DataFrame = None) -> tuple:
    """生成与汇率相关的政策数据"""
    date_range = pd.date_range(start=start_date, end=end_date)
    
    # 动态生成政策事件（基于汇率波动）
    policies = {}
    
    if exchange_rate_data is not None and len(exchange_rate_data) > 0:
        # 基于汇率波动生成政策事件
        rate_volatility = exchange_rate_data['exchange_rate'].rolling(window=10).std().fillna(0)
        rate_changes = exchange_rate_data['exchange_rate'].pct_change().fillna(0)
        
        # 在高波动期或大幅变化时生成政策事件
        high_volatility_dates = rate_volatility[rate_volatility > rate_volatility.quantile(0.8)].index
        large_change_dates = rate_changes[abs(rate_changes) > rate_changes.std() * 2].index
        
        policy_dates = list(set(list(high_volatility_dates) + list(large_change_dates)))
        
        policy_types = [
            {'name': '央行降准', 'base_impact': -0.4},
            {'name': '利率调整', 'base_impact': 0.3},
            {'name': '汇率政策', 'base_impact': -0.2},
            {'name': '货币政策', 'base_impact': 0.35},
            {'name': '外汇管制', 'base_impact': -0.3},
            {'name': '市场干预', 'base_impact': 0.25}
        ]
        
        # 选择部分日期作为政策事件日
        selected_dates = np.random.choice(policy_dates, 
                                        size=min(6, len(policy_dates)), 
                                        replace=False) if len(policy_dates) > 0 else []
        
        for i, date in enumerate(selected_dates):
            policy_type = policy_types[i % len(policy_types)]
            # 根据当时的汇率变化调整影响强度
            rate_change = rate_changes.loc[date] if date in rate_changes.index else 0
            impact_adjustment = -0.5 * rate_change  # 政策通常与汇率变化方向相反
            
            policies[date.strftime('%Y-%m-%d')] = {
                'name': policy_type['name'],
                'impact': policy_type['base_impact'] + impact_adjustment,
                'description': f"{policy_type['name']}响应市场波动"
            }
    else:
        # 默认政策事件
        policies = {
            '2023-03-15': {'name': '央行降准', 'impact': -0.5, 'description': '央行宣布降准0.5个百分点'},
            '2023-06-20': {'name': '利率调整', 'impact': 0.3, 'description': '央行上调基准利率'},
            '2023-09-10': {'name': '汇率政策', 'impact': -0.2, 'description': '外汇管理政策调整'},
            '2023-11-05': {'name': '货币政策', 'impact': 0.4, 'description': '货币政策委员会会议'}
        }
    
    # 创建政策影响时间序列
    policy_impact = np.zeros(len(date_range))
    
    for date_str, policy in policies.items():
        policy_date = pd.to_datetime(date_str)
        if policy_date in date_range:
            idx = date_range.get_loc(policy_date)
            impact = policy['impact']
            
            # 影响衰减
            for i in range(idx, min(idx + 30, len(date_range))):
                decay = np.exp(-(i - idx) / 10)
                policy_impact[i] += impact * decay
    
    policy_df = pd.DataFrame({
        'policy_impact': policy_impact
    }, index=date_range)
    
    return policy_df, policies

def generate_event_data(start_date: str, end_date: str, exchange_rate_data: pd.DataFrame = None) -> tuple:
    """生成与汇率相关的重大事件数据"""
    date_range = pd.date_range(start=start_date, end=end_date)
    
    # 动态生成重大事件（基于汇率异常波动）
    events = {}
    
    if exchange_rate_data is not None and len(exchange_rate_data) > 0:
        # 基于汇率异常波动生成事件
        rate_changes = exchange_rate_data['exchange_rate'].pct_change().fillna(0)
        rate_volatility = exchange_rate_data['exchange_rate'].rolling(window=5).std().fillna(0)
        
        # 识别异常波动日期
        extreme_change_dates = rate_changes[abs(rate_changes) > rate_changes.std() * 2.5].index
        high_volatility_dates = rate_volatility[rate_volatility > rate_volatility.quantile(0.9)].index
        
        event_dates = list(set(list(extreme_change_dates) + list(high_volatility_dates)))
        
        event_types = [
            {'name': '地缘政治事件', 'base_impact': 0.6},
            {'name': '经济数据发布', 'base_impact': -0.3},
            {'name': '贸易协议', 'base_impact': -0.4},
            {'name': '金融危机', 'base_impact': 0.8},
            {'name': '央行会议', 'base_impact': 0.4},
            {'name': '国际制裁', 'base_impact': 0.7},
            {'name': '经济刺激', 'base_impact': -0.5}
        ]
        
        # 选择部分日期作为重大事件日
        selected_dates = np.random.choice(event_dates, 
                                        size=min(5, len(event_dates)), 
                                        replace=False) if len(event_dates) > 0 else []
        
        for i, date in enumerate(selected_dates):
            event_type = event_types[i % len(event_types)]
            # 根据当时的汇率变化调整影响强度
            rate_change = rate_changes.loc[date] if date in rate_changes.index else 0
            impact_multiplier = 1 + abs(rate_change) * 10  # 变化越大，事件影响越强
            
            events[date.strftime('%Y-%m-%d')] = {
                'name': event_type['name'],
                'impact': event_type['base_impact'] * impact_multiplier,
                'description': f"{event_type['name']}引发市场剧烈波动"
            }
    else:
        # 默认重大事件
        events = {
            '2023-02-10': {'name': '地缘政治事件', 'impact': 0.6, 'description': '国际地缘政治紧张局势'},
            '2023-05-15': {'name': '经济数据发布', 'impact': -0.3, 'description': 'GDP数据超预期'},
            '2023-08-20': {'name': '贸易协议', 'impact': -0.4, 'description': '重要贸易协议签署'},
            '2023-10-30': {'name': '金融危机', 'impact': 0.8, 'description': '区域金融市场波动'}
        }
    
    # 创建事件影响时间序列
    event_impact = np.zeros(len(date_range))
    
    for date_str, event in events.items():
        event_date = pd.to_datetime(date_str)
        if event_date in date_range:
            idx = date_range.get_loc(event_date)
            impact = event['impact']
            
            # 影响衰减
            for i in range(idx, min(idx + 20, len(date_range))):
                decay = np.exp(-(i - idx) / 7)
                event_impact[i] += impact * decay
    
    event_df = pd.DataFrame({
        'event_impact': event_impact
    }, index=date_range)
    
    return event_df, events

def analyze_factor_contribution(combined_df: pd.DataFrame) -> Dict:
    """分析各因素对汇率的贡献度 - 优化版本"""
    # 创建更丰富的特征集
    df = combined_df.copy()
    
    # 添加技术指标特征
    df['exchange_rate_ma5'] = df['exchange_rate'].rolling(window=5).mean()
    df['exchange_rate_ma10'] = df['exchange_rate'].rolling(window=10).mean()
    df['exchange_rate_volatility'] = df['exchange_rate'].rolling(window=5).std()
    df['exchange_rate_change'] = df['exchange_rate'].pct_change()
    df['exchange_rate_lag1'] = df['exchange_rate'].shift(1)
    df['exchange_rate_lag2'] = df['exchange_rate'].shift(2)
    
    # 添加交互特征
    df['sentiment_policy_interaction'] = df['sentiment'] * df['policy_impact']
    df['sentiment_event_interaction'] = df['sentiment'] * df['event_impact']
    df['policy_event_interaction'] = df['policy_impact'] * df['event_impact']
    
    # 添加滞后的情绪和政策影响
    df['sentiment_lag1'] = df['sentiment'].shift(1)
    df['policy_impact_lag1'] = df['policy_impact'].shift(1)
    df['event_impact_lag1'] = df['event_impact'].shift(1)
    
    # 添加累积影响
    df['sentiment_cumsum'] = df['sentiment'].rolling(window=7).mean()
    df['policy_cumsum'] = df['policy_impact'].rolling(window=7).sum()
    df['event_cumsum'] = df['event_impact'].rolling(window=7).sum()
    
    # 删除包含NaN的行
    df = df.dropna()
    
    if len(df) < 10:  # 确保有足够的数据
        print("警告：数据量不足，使用简化模型")
        return {
            'sentiment_contrib': 0.33,
            'policy_contrib': 0.33,
            'event_contrib': 0.34,
            'r2': 0.0253
        }
    
    # 准备特征和目标变量
    feature_columns = [
        'sentiment', 'policy_impact', 'event_impact',
        'exchange_rate_ma5', 'exchange_rate_ma10', 'exchange_rate_volatility',
        'exchange_rate_change', 'exchange_rate_lag1', 'exchange_rate_lag2',
        'sentiment_policy_interaction', 'sentiment_event_interaction', 'policy_event_interaction',
        'sentiment_lag1', 'policy_impact_lag1', 'event_impact_lag1',
        'sentiment_cumsum', 'policy_cumsum', 'event_cumsum'
    ]
    
    X = df[feature_columns].copy()
    y = df['exchange_rate'].copy()
    
    # 填充缺失值
    X = X.fillna(X.mean())
    y = y.fillna(method='ffill')
    
    # 标准化特征
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    # 训练线性回归模型
    model = LinearRegression()
    model.fit(X_scaled, y)
    
    # 计算R²
    r2 = model.score(X_scaled, y)
    
    # 计算特征重要性
    feature_importance = np.abs(model.coef_)
    
    # 计算原始三个因素的贡献度（基于相关特征的重要性加权）
    sentiment_indices = [0, 9, 12, 15]  # sentiment相关特征的索引
    policy_indices = [1, 10, 13, 16]    # policy相关特征的索引
    event_indices = [2, 11, 14, 17]     # event相关特征的索引
    
    sentiment_importance = np.sum(feature_importance[sentiment_indices])
    policy_importance = np.sum(feature_importance[policy_indices])
    event_importance = np.sum(feature_importance[event_indices])
    
    total_importance = sentiment_importance + policy_importance + event_importance
    
    if total_importance == 0:
        contrib = np.array([1/3, 1/3, 1/3])
    else:
        contrib = np.array([sentiment_importance, policy_importance, event_importance]) / total_importance
    
    print(f"优化后模型性能提升：")
    print(f"  使用特征数量: {len(feature_columns)}")
    print(f"  训练样本数量: {len(X)}")
    print(f"  模型拟合优度(R²): {r2:.4f}")
    
    return {
        'sentiment_contrib': contrib[0],
        'policy_contrib': contrib[1],
        'event_contrib': contrib[2],
        'r2': r2,
        'feature_count': len(feature_columns),
        'sample_count': len(X)
    }

def create_comprehensive_analysis_plot(combined_df: pd.DataFrame, policies: Dict, 
                                     events: Dict, currency_pair: str, 
                                     analysis_results: Dict) -> str:
    """创建综合分析图表"""
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(20, 15))
    
    # 1. 汇率与情绪分析
    ax1_twin = ax1.twinx()
    ax1.plot(combined_df.index, combined_df['exchange_rate'], 'b-', label=f'{currency_pair}汇率')
    ax1_twin.plot(combined_df.index, combined_df['sentiment'], 'r-', label='舆论情绪')
    ax1.set_ylabel('汇率', color='b')
    ax1_twin.set_ylabel('情绪指数', color='r')
    ax1.set_title('汇率与舆论情绪关系')
    ax1.grid(True, alpha=0.3)
    
    # 2. 汇率与政策影响
    ax2_twin = ax2.twinx()
    ax2.plot(combined_df.index, combined_df['exchange_rate'], 'b-', label=f'{currency_pair}汇率')
    ax2_twin.plot(combined_df.index, combined_df['policy_impact'], 'g-', label='政策影响')
    ax2.set_ylabel('汇率', color='b')
    ax2_twin.set_ylabel('政策影响', color='g')
    ax2.set_title('汇率与政策影响关系')
    ax2.grid(True, alpha=0.3)
    
    # 标记政策事件
    for date_str, policy in policies.items():
        policy_date = pd.to_datetime(date_str)
        if policy_date in combined_df.index:
            ax2.axvline(x=policy_date, color='gray', linestyle='--', alpha=0.7)
    
    # 3. 汇率与事件影响
    ax3_twin = ax3.twinx()
    ax3.plot(combined_df.index, combined_df['exchange_rate'], 'b-', label=f'{currency_pair}汇率')
    ax3_twin.plot(combined_df.index, combined_df['event_impact'], 'm-', label='事件影响')
    ax3.set_ylabel('汇率', color='b')
    ax3_twin.set_ylabel('事件影响', color='m')
    ax3.set_title('汇率与重大事件关系')
    ax3.grid(True, alpha=0.3)
    
    # 标记重大事件
    for date_str, event in events.items():
        event_date = pd.to_datetime(date_str)
        if event_date in combined_df.index:
            ax3.axvline(x=event_date, color='gray', linestyle='--', alpha=0.7)
    
    # 4. 因素贡献度饼图
    labels = ['舆论情绪', '政策影响', '事件影响']
    sizes = [analysis_results['sentiment_contrib'], 
            analysis_results['policy_contrib'], 
            analysis_results['event_contrib']]
    colors = ['#ff9999','#66b3ff','#99ff99']
    
    ax4.pie(sizes, labels=labels, colors=colors, autopct='%1.1f%%', startangle=90)
    ax4.set_title('影响因素贡献度分析')
    
    plt.tight_layout()
    
    # 保存图表
    filename = f"{currency_pair.replace('/', '_')}_comprehensive_analysis.png"
    filepath = os.path.join(output_dir, filename)
    plt.savefig(filepath, dpi=300, bbox_inches='tight')
    plt.close()
    
    return filepath

# --- 4. LangGraph Agent节点定义 ---
def initialize_system_node(state: MultiAgentState) -> MultiAgentState:
    """系统初始化节点"""
    print("=== 多Agent汇率预测分析系统启动 ===")
    state['current_agent'] = 'system_initializer'
    state['agent_results'] = {}
    
    # 设置默认值
    if not state.get('currency_pair'):
        state['currency_pair'] = 'USD/CNY'
    if not state.get('start_date'):
        state['start_date'] = '2023-01-01'
    if not state.get('end_date'):
        state['end_date'] = '2023-12-31'
    
    print(f"分析货币对: {state['currency_pair']}")
    print(f"分析时间范围: {state['start_date']} 到 {state['end_date']}")
    
    state['next_action'] = 'data_collection'
    state['error_message'] = None
    return state

def data_collection_agent_node(state: MultiAgentState) -> MultiAgentState:
    """数据收集Agent节点 - 优化版本"""
    print("\n[数据收集Agent] 开始收集数据...")
    state['current_agent'] = 'data_collector'
    
    try:
        # 首先加载汇率数据
        state['raw_data'] = load_exchange_rate_data(
            currency_pair=state['currency_pair'],
            start_date=state['start_date'],
            end_date=state['end_date']
        )
        
        print(f"汇率数据加载完成: {len(state['raw_data'])}条记录")
        
        # 基于汇率数据生成相关的舆论情绪数据
        state['sentiment_data'] = generate_sentiment_data(
            state['start_date'], 
            state['end_date'], 
            exchange_rate_data=state['raw_data']
        )
        
        # 基于汇率数据生成相关的政策数据
        state['policy_data'], state['policy_events'] = generate_policy_data(
            state['start_date'], 
            state['end_date'], 
            exchange_rate_data=state['raw_data']
        )
        
        # 基于汇率数据生成相关的事件数据
        state['event_data'], state['major_events'] = generate_event_data(
            state['start_date'], 
            state['end_date'], 
            exchange_rate_data=state['raw_data']
        )
        
        print(f"数据收集完成 - 汇率数据: {len(state['raw_data'])}条")
        print(f"情绪数据: {len(state['sentiment_data'])}条")
        print(f"政策事件: {len(state['policy_events'])}个")
        print(f"重大事件: {len(state['major_events'])}个")
        print("注意：情绪、政策和事件数据已基于汇率波动生成，提高了相关性")
        
        state['agent_results']['data_collection'] = {
            'status': 'success',
            'data_points': len(state['raw_data']),
            'events_count': len(state['policy_events']) + len(state['major_events']),
            'correlation_enhanced': True
        }
        
        state['next_action'] = 'model_training'
        state['error_message'] = None
        
    except Exception as e:
        print(f"数据收集失败: {str(e)}")
        state['error_message'] = f"数据收集失败: {str(e)}"
        state['next_action'] = 'handle_error'
    
    return state

def model_training_agent_node(state: MultiAgentState) -> MultiAgentState:
    """模型训练Agent节点"""
    print("\n[模型训练Agent] 开始训练预测模型...")
    state['current_agent'] = 'model_trainer'
    
    try:
        # 创建特征
        X, y, dates = create_features(state['raw_data'])
        
        # 训练模型
        state['model'], state['scaler'] = train_model(X, y)
        
        print("模型训练完成")
        
        state['agent_results']['model_training'] = {
            'status': 'success',
            'features_count': len(X.columns),
            'training_samples': len(X)
        }
        
        state['next_action'] = 'sentiment_analysis'
        state['error_message'] = None
        
    except Exception as e:
        print(f"模型训练失败: {str(e)}")
        state['error_message'] = f"模型训练失败: {str(e)}"
        state['next_action'] = 'handle_error'
    
    return state

def sentiment_analysis_agent_node(state: MultiAgentState) -> MultiAgentState:
    """情绪分析Agent节点"""
    print("\n[情绪分析Agent] 开始分析舆论情绪、政策和事件影响...")
    state['current_agent'] = 'sentiment_analyzer'
    
    try:
        # 创建综合分析数据
        combined_df = pd.DataFrame({
            'exchange_rate': state['raw_data']['exchange_rate'],
            'sentiment': state['sentiment_data']['sentiment'],
            'policy_impact': state['policy_data']['policy_impact'],
            'event_impact': state['event_data']['event_impact']
        })
        
        state['combined_analysis'] = combined_df
        
        # 分析各因素贡献度
        state['analysis_results'] = analyze_factor_contribution(combined_df)
        
        print(f"情绪分析完成:")
        print(f"  舆论情绪贡献度: {state['analysis_results']['sentiment_contrib']:.2%}")
        print(f"  政策影响贡献度: {state['analysis_results']['policy_contrib']:.2%}")
        print(f"  事件影响贡献度: {state['analysis_results']['event_contrib']:.2%}")
        print(f"  模型拟合优度(R²): {state['analysis_results']['r2']:.4f}")
        
        state['agent_results']['sentiment_analysis'] = {
            'status': 'success',
            'r2_score': state['analysis_results']['r2'],
            'sentiment_contrib': state['analysis_results']['sentiment_contrib'],
            'policy_contrib': state['analysis_results']['policy_contrib'],
            'event_contrib': state['analysis_results']['event_contrib']
        }
        
        state['next_action'] = 'get_user_query'
        state['error_message'] = None
        
    except Exception as e:
        print(f"情绪分析失败: {str(e)}")
        state['error_message'] = f"情绪分析失败: {str(e)}"
        state['next_action'] = 'handle_error'
    
    return state

def user_interaction_agent_node(state: MultiAgentState) -> MultiAgentState:
    """用户交互Agent节点"""
    print("\n[用户交互Agent] 等待用户输入...")
    state['current_agent'] = 'user_interface'
    
    try:
        print("\n=== 汇率预测选项 ===")
        print("1. 预测未来汇率")
        print("2. 查看综合分析报告")
        print("3. 退出系统")
        
        choice = input("请选择操作 (1-3): ").strip()
        
        if choice == '1':
            days = input("请输入预测天数 (默认30天): ").strip()
            try:
                state['prediction_days'] = int(days) if days else 30
            except ValueError:
                state['prediction_days'] = 30
            
            state['user_query'] = 'predict'
            state['next_action'] = 'prediction'
            
        elif choice == '2':
            state['user_query'] = 'report'
            state['next_action'] = 'generate_report'
            
        elif choice == '3':
            state['user_query'] = 'exit'
            state['next_action'] = 'end_conversation'
            
        else:
            print("无效选择，请重新输入")
            state['next_action'] = 'get_user_query'
        
        state['error_message'] = None
        
    except Exception as e:
        print(f"用户交互失败: {str(e)}")
        state['error_message'] = f"用户交互失败: {str(e)}"
        state['next_action'] = 'handle_error'
    
    return state

def prediction_agent_node(state: MultiAgentState) -> MultiAgentState:
    """预测Agent节点"""
    print(f"\n[预测Agent] 开始预测未来{state['prediction_days']}天的汇率...")
    state['current_agent'] = 'predictor'
    
    try:
        # 获取最近的数据用于预测
        recent_data = state['raw_data'].tail(50)  # 使用最近50天的数据
        
        # 预测未来汇率
        predictions, prediction_dates = predict_future(
            state['model'], 
            state['scaler'], 
            recent_data, 
            state['prediction_days']
        )
        
        state['predictions'] = predictions
        state['prediction_dates'] = prediction_dates
        
        print(f"预测完成，未来{len(predictions)}天的汇率趋势:")
        print(f"  起始预测值: {predictions[0]:.4f}")
        print(f"  结束预测值: {predictions[-1]:.4f}")
        print(f"  预测变化: {((predictions[-1] - predictions[0]) / predictions[0] * 100):+.2f}%")
        
        # 显示详细的预测数值
        print("\n详细预测数值:")
        if len(predictions) <= 15:
            # 如果预测天数不多，显示所有数值
            for i, (date, value) in enumerate(zip(prediction_dates, predictions)):
                print(f"  {date.strftime('%Y-%m-%d')}: {value:.4f}")
        else:
            # 如果预测天数较多，显示前10天和后5天
            print("  前10天预测:")
            for i in range(10):
                date = prediction_dates[i]
                value = predictions[i]
                print(f"    {date.strftime('%Y-%m-%d')}: {value:.4f}")
            
            print(f"  ... (省略{len(predictions)-15}天) ...")
            
            print("  后5天预测:")
            for i in range(-5, 0):
                date = prediction_dates[i]
                value = predictions[i]
                print(f"    {date.strftime('%Y-%m-%d')}: {value:.4f}")
        
        state['agent_results']['prediction'] = {
            'status': 'success',
            'prediction_days': len(predictions),
            'start_value': float(predictions[0]),
            'end_value': float(predictions[-1]),
            'change_percent': float((predictions[-1] - predictions[0]) / predictions[0] * 100)
        }
        
        state['next_action'] = 'generate_visualization'
        state['error_message'] = None
        
    except Exception as e:
        print(f"预测失败: {str(e)}")
        state['error_message'] = f"预测失败: {str(e)}"
        state['next_action'] = 'handle_error'
    
    return state

def visualization_agent_node(state: MultiAgentState) -> MultiAgentState:
    """可视化Agent节点"""
    print("\n[可视化Agent] 生成图表和报告...")
    state['current_agent'] = 'visualizer'
    
    try:
        # 生成预测图表
        if state.get('predictions') is not None and state.get('prediction_dates') is not None:
            forecast_plot = plot_forecast(
                state['raw_data'], 
                state['prediction_dates'], 
                state['predictions'], 
                state['currency_pair']
            )
            print(f"预测图表已保存: {forecast_plot}")
        
        # 生成综合分析图表
        if state.get('combined_analysis') is not None:
            analysis_plot = create_comprehensive_analysis_plot(
                state['combined_analysis'],
                state['policy_events'],
                state['major_events'],
                state['currency_pair'],
                state['analysis_results']
            )
            print(f"综合分析图表已保存: {analysis_plot}")
        
        state['agent_results']['visualization'] = {
            'status': 'success',
            'charts_generated': 2
        }
        
        state['next_action'] = 'get_user_query'
        state['error_message'] = None
        
    except Exception as e:
        print(f"可视化生成失败: {str(e)}")
        state['error_message'] = f"可视化生成失败: {str(e)}"
        state['next_action'] = 'handle_error'
    
    return state

def report_generation_agent_node(state: MultiAgentState) -> MultiAgentState:
    """报告生成Agent节点"""
    print("\n[报告生成Agent] 生成综合分析报告...")
    state['current_agent'] = 'report_generator'
    
    try:
        print("\n=== 综合分析报告 ===")
        print(f"货币对: {state['currency_pair']}")
        print(f"分析期间: {state['start_date']} 至 {state['end_date']}")
        print(f"数据点数: {len(state['raw_data'])}")
        
        print("\n--- 影响因素分析 ---")
        print(f"舆论情绪贡献度: {state['analysis_results']['sentiment_contrib']:.2%}")
        print(f"政策影响贡献度: {state['analysis_results']['policy_contrib']:.2%}")
        print(f"事件影响贡献度: {state['analysis_results']['event_contrib']:.2%}")
        print(f"模型拟合优度(R²): {state['analysis_results']['r2']:.4f}")
        
        print("\n--- 政策事件 ---")
        for date, policy in state['policy_events'].items():
            print(f"{date}: {policy['name']} (影响: {policy['impact']:+.2f})")
        
        print("\n--- 重大事件 ---")
        for date, event in state['major_events'].items():
            print(f"{date}: {event['name']} (影响: {event['impact']:+.2f})")
        
        # 生成综合分析图表
        analysis_plot = create_comprehensive_analysis_plot(
            state['combined_analysis'],
            state['policy_events'],
            state['major_events'],
            state['currency_pair'],
            state['analysis_results']
        )
        print(f"\n综合分析图表已保存: {analysis_plot}")
        
        state['agent_results']['report_generation'] = {
            'status': 'success',
            'report_generated': True
        }
        
        state['next_action'] = 'get_user_query'
        state['error_message'] = None
        
    except Exception as e:
        print(f"报告生成失败: {str(e)}")
        state['error_message'] = f"报告生成失败: {str(e)}"
        state['next_action'] = 'handle_error'
    
    return state

def error_handling_agent_node(state: MultiAgentState) -> MultiAgentState:
    """错误处理Agent节点"""
    print(f"\n[错误处理Agent] 处理错误: {state.get('error_message', '未知错误')}")
    state['current_agent'] = 'error_handler'
    
    print("尝试恢复系统...")
    
    # 简单的错误恢复策略
    if 'data_collection' in state.get('agent_results', {}):
        state['next_action'] = 'get_user_query'
    else:
        state['next_action'] = 'data_collection'
    
    return state

def end_conversation_node(state: MultiAgentState) -> MultiAgentState:
    """结束对话节点"""
    print("\n=== 系统结束 ===")
    print("感谢使用多Agent汇率预测分析系统！")
    
    # 打印系统运行总结
    print("\n--- 系统运行总结 ---")
    for agent, results in state.get('agent_results', {}).items():
        print(f"{agent}: {results.get('status', 'unknown')}")
    
    return state

# --- 5. 构建LangGraph工作流 ---
def create_multi_agent_workflow() -> StateGraph:
    """创建多Agent工作流"""
    workflow = StateGraph(MultiAgentState)
    
    # 添加Agent节点
    workflow.add_node("initialize_system", initialize_system_node)
    workflow.add_node("data_collection", data_collection_agent_node)
    workflow.add_node("model_training", model_training_agent_node)
    workflow.add_node("sentiment_analysis", sentiment_analysis_agent_node)
    workflow.add_node("get_user_query", user_interaction_agent_node)
    workflow.add_node("prediction", prediction_agent_node)
    workflow.add_node("generate_visualization", visualization_agent_node)
    workflow.add_node("generate_report", report_generation_agent_node)
    workflow.add_node("handle_error", error_handling_agent_node)
    workflow.add_node("end_conversation", end_conversation_node)
    
    # 设置入口点
    workflow.set_entry_point("initialize_system")
    
    # 定义路由逻辑
    def route_next_action(state: MultiAgentState) -> str:
        return state.get('next_action', 'end_conversation')
    
    # 添加条件边
    workflow.add_conditional_edges(
        "initialize_system",
        route_next_action,
        {
            "data_collection": "data_collection",
            "end_conversation": END
        }
    )
    
    workflow.add_conditional_edges(
        "data_collection",
        route_next_action,
        {
            "model_training": "model_training",
            "handle_error": "handle_error",
            "end_conversation": END
        }
    )
    
    workflow.add_conditional_edges(
        "model_training",
        route_next_action,
        {
            "sentiment_analysis": "sentiment_analysis",
            "handle_error": "handle_error",
            "end_conversation": END
        }
    )
    
    workflow.add_conditional_edges(
        "sentiment_analysis",
        route_next_action,
        {
            "get_user_query": "get_user_query",
            "handle_error": "handle_error",
            "end_conversation": END
        }
    )
    
    workflow.add_conditional_edges(
        "get_user_query",
        route_next_action,
        {
            "prediction": "prediction",
            "generate_report": "generate_report",
            "get_user_query": "get_user_query",
            "handle_error": "handle_error",
            "end_conversation": END
        }
    )
    
    workflow.add_conditional_edges(
        "prediction",
        route_next_action,
        {
            "generate_visualization": "generate_visualization",
            "handle_error": "handle_error",
            "end_conversation": END
        }
    )
    
    workflow.add_conditional_edges(
        "generate_visualization",
        route_next_action,
        {
            "get_user_query": "get_user_query",
            "handle_error": "handle_error",
            "end_conversation": END
        }
    )
    
    workflow.add_conditional_edges(
        "generate_report",
        route_next_action,
        {
            "get_user_query": "get_user_query",
            "handle_error": "handle_error",
            "end_conversation": END
        }
    )
    
    workflow.add_conditional_edges(
        "handle_error",
        route_next_action,
        {
            "data_collection": "data_collection",
            "get_user_query": "get_user_query",
            "end_conversation": END
        }
    )
    
    return workflow

# --- 6. 主程序 ---
def main():
    """主程序入口"""
    print("启动多Agent汇率预测分析系统...")
    
    # 创建工作流
    workflow = create_multi_agent_workflow()
    app = workflow.compile()
    
    # 获取当前日期，设置历史数据范围为过去一年到今天
    from datetime import datetime, timedelta
    today = datetime.now()
    one_year_ago = today - timedelta(days=365)
    
    # 初始状态
    initial_state = {
        'currency_pair': 'USD/CNY',
        'start_date': one_year_ago.strftime('%Y-%m-%d'),
        'end_date': today.strftime('%Y-%m-%d'),
        'next_action': 'data_collection'
    }
    
    # 运行系统
    try:
        current_state = initial_state
        while True:
            current_state = app.invoke(current_state)
            
            # 检查是否结束
            if current_state.get('next_action') == 'end_conversation':
                break
            
            # 如果是用户交互节点，需要重新调用以获取用户输入
            if current_state.get('current_agent') == 'user_interface':
                continue
                
    except KeyboardInterrupt:
        print("\n用户中断，系统退出")
    except Exception as e:
        print(f"\n系统运行错误: {str(e)}")
    
    print("\n系统已退出")

if __name__ == '__main__':
    main()